import fs from "fs";
import { PDFDocument } from "pdf-lib";

export async function mergePDFsToA4Grid(
  pdfPaths: string[],
  outputPath: string,
  columns = 1,
  rows = 2
): Promise<string> {
  const A4_WIDTH = 595.28;
  const A4_HEIGHT = 841.89;
  const margin = 10;
  const cellWidth = (A4_WIDTH - margin * (columns + 1)) / columns;
  const cellHeight = (A4_HEIGHT - margin * (rows + 1)) / rows;

  const newPdf = await PDFDocument.create();
  let page = newPdf.addPage([A4_WIDTH, A4_HEIGHT]);
  let index = 0;

  for (let i = 0; i < pdfPaths.length; i++) {
    const bytes = fs.readFileSync(pdfPaths[i]);
    const srcPdf = await PDFDocument.load(bytes);
    const [srcPage] = await srcPdf.getPages(); // ⭐ 来源 PDF 的第一页

    const { width, height } = srcPage.getSize();

    // ⭐ 必须先“嵌入”该页面，否则 drawPage 会报错
    const embeddedPage = await newPdf.embedPage(srcPage);

    const scale = Math.min(cellWidth / width, cellHeight / height);
    const scaledWidth = width * scale;
    const scaledHeight = height * scale;

    const col = index % columns;
    const row = Math.floor(index / columns) % rows;

    const x =
      margin + col * (cellWidth + margin) + (cellWidth - scaledWidth) / 2;
    const y =
      A4_HEIGHT -
      margin -
      (row + 1) * (cellHeight + margin) +
      (cellHeight - scaledHeight) / 2;

    // ⭐ 使用 drawPage 绘制“嵌入的页面”
    page.drawPage(embeddedPage, {
      x,
      y,
      xScale: scale,
      yScale: scale,
    });

    index++;

    if (index % (columns * rows) === 0 && i !== pdfPaths.length - 1) {
      page = newPdf.addPage([A4_WIDTH, A4_HEIGHT]);
    }
  }

  const mergedBytes = await newPdf.save();
  fs.writeFileSync(outputPath, mergedBytes);
  return outputPath;
}
